from typing import TypedDict
import uuid

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.constants import START
from langgraph.graph import StateGraph
from langgraph.types import interrupt, Command

class State(TypedDict):
    some_text: str
    revision_count: int

def human_node(state: State):
    """人工审核节点 - 中断执行等待人工输入"""
    print("\n*********进入人工审核节点*********")
    print(f"\n当前文本: '{state['some_text']}'")
    print(f"修订次数: {state.get('revision_count', 0)}")
    
    # 中断并等待人工输入
    value = interrupt({
        "text_to_revise": state["some_text"],
        "message": "请提供修改后的文本，或输入 'ok' 接受当前文本"
    })
    
    # 更新修订计数
    revision_count = state.get('revision_count', 0)
    if value != state['some_text'] and value != 'ok':
        revision_count += 1
    
    return {
        "some_text": value if value != 'ok' else state['some_text'],
        "revision_count": revision_count
    }

# 构建图
graph_builder = StateGraph(State)
graph_builder.add_node("human_node", human_node)
graph_builder.add_edge(START, "human_node")
checkpointer = InMemorySaver()
graph = graph_builder.compile(checkpointer=checkpointer)

# 演示函数
def demo_hitl():
    config = {"configurable": {"thread_id": str(uuid.uuid4())}}
    
    # 初始文本
    initial_state = {"some_text": "原始文本", "revision_count": 0}
    
    print("=== Human-in-the-Loop 演示 ===")
    
    # 第一次运行 - 会在中断处停止
    result = graph.invoke(initial_state, config)
    
    interrupt_info = result['__interrupt__'][0].value
    
    # 模拟人工输入
    print(f"\n当前文本：{interrupt_info['text_to_revise']}")
    user_input = input(f"{interrupt_info['message']}: ")
    
    # 恢复执行
    final_result = graph.invoke(Command(resume=user_input), config)
    
    print(f"\n最终结果:")
    print(f"文本: '{final_result['some_text']}'")
    print(f"修订次数: {final_result['revision_count']}")

if __name__ == "__main__":
    demo_hitl()